from read_CloudSat import reader
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from datetime import datetime
import pandas as pd
import netCDF4 as nc
from scipy.spatial import cKDTree


def set_map(ax):
    '''设置全球地图背景.'''
    proj = ccrs.PlateCarree()
    ax.coastlines(lw=0.5)
    xticks = np.arange(-180, 180 + 60, 60)
    yticks = np.arange(-90, 90 + 30, 30)
    ax.set_xticks(xticks, crs=proj)
    ax.set_yticks(yticks, crs=proj)
    ax.xaxis.set_major_formatter(LongitudeFormatter())
    ax.yaxis.set_major_formatter(LatitudeFormatter())
    ax.tick_params('both', labelsize='x-small')
    ax.set_global()


def draw_track(ax, lon1D, lat1D, lat_range=None):
    '''
    根据经纬度画出轨迹,并标识出起点.
    新增功能：可根据纬度范围筛选轨迹段进行绘制
    '''
    # 如果提供了纬度范围，则筛选该范围内的轨迹
    if lat_range:
        mask = (lat1D >= lat_range[0]) & (lat1D <= lat_range[1])
        if np.any(mask):
            lon_segment = lon1D[mask]
            lat_segment = lat1D[mask]
            # 绘制筛选后的轨迹段
            ax.plot(lon_segment, lat_segment, lw=2, color='b', transform=ccrs.Geodetic())
            # 绘制筛选后轨迹的起点
            ax.plot(lon_segment[0], lat_segment[0], 'ro', ms=3, transform=ccrs.PlateCarree())
            ax.text(
                lon_segment[0] + 5, lat_segment[0], 'start', color='r', fontsize='x-small',
                transform=ccrs.PlateCarree()
            )
    else:
        # 未提供纬度范围时绘制完整轨迹
        ax.plot(lon1D, lat1D, lw=2, color='b', transform=ccrs.Geodetic())
        ax.plot(lon1D[0], lat1D[0], 'ro', ms=3, transform=ccrs.PlateCarree())
        ax.text(
            lon1D[0] + 5, lat1D[0], 'start', color='r', fontsize='x-small',
            transform=ccrs.PlateCarree()
        )


# ------------------ MODIFICATION: Add color and linewidth parameters ------------------
def draw_highlight_box(ax, lon_full_track, lat_full_track, highlight_lat_range,
                       box_color='red', box_linewidth=2):
    '''
    在地图上绘制一个高亮框，以显示特定纬度范围内的轨迹段。

    参数:
    ax: Matplotlib aexs 对象.
    lon_full_track, lat_full_track: 完整的轨迹经纬度数据.
    highlight_lat_range: 要高亮的纬度范围, e.g., (15, 33).
    box_color (str): 框的颜色.
    box_linewidth (int or float): 框的线宽.
    '''
    mask = (lat_full_track >= highlight_lat_range[0]) & (lat_full_track <= highlight_lat_range[1])
    if not np.any(mask):
        return

    lon_segment = lon_full_track[mask]
    lat_segment = lat_full_track[mask]

    min_lon, max_lon = lon_segment.min(), lon_segment.max()
    min_lat, max_lat = lat_segment.min(), lat_segment.max()

    box_lons = [min_lon, max_lon, max_lon, min_lon, min_lon]
    box_lats = [min_lat, min_lat, max_lat, max_lat, min_lat]

    # 使用传入的参数来设置颜色和线宽
    ax.plot(box_lons, box_lats, color=box_color, linestyle='-',
            linewidth=box_linewidth, transform=ccrs.PlateCarree())


def draw_cross_section(ax, lon, hgt, data, cth_sorted, cbh_sorted, latitude_sorted, xlims=None):
    '''画出data的经度-高度剖面，并添加风云数据的云底和云顶高度折线图.'''
    im = ax.pcolormesh(lon, hgt, data, cmap='jet', shading='nearest')
    ax.set_ylim(0, 20)

    if xlims:
        ax.set_xlim(xlims)
    else:
        ax.set_xlim(lon.min(), lon.max())

    ax.set_xlabel('Latitude', fontsize='x-small')
    ax.set_ylabel('Height [km]', fontsize='x-small')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('bottom', size='10%', pad=0.42)
    cbar = plt.colorbar(im, cax=cax, extend='both', orientation='horizontal')
    cbar.ax.tick_params(labelsize='x-small')
    cbar.set_label('dBZe', fontsize='x-small')

    ax.plot(latitude_sorted, cth_sorted, color='black', label='CTH vs Longitude', linewidth=0.8)
    ax.plot(latitude_sorted, cbh_sorted, color='red', label='CBH vs Longitude', linewidth=0.8)
    ax.legend(fontsize='x-small')


def draw_elevation(ax, lon, elv):
    '''画出地形抬升.'''
    ax.fill_between(lon, elv, color='gray')


if __name__ == '__main__':
    # (数据读取和处理部分保持不变)
    # ...
    fname = (
        # '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020139071017_74881_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
        # '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020159081830_75174_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
        '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020142040823_74923_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
        # '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506/2020165053142_75260_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
    )
    f = reader(fname)
    lon, lat, elv = f.read_geo()
    data = f.read_sds('Radar_Reflectivity')
    height = f.read_sds('Height')
    time = f.read_time(datetime=True)
    f.close()

    start_datetime = datetime(2020, 5, 21, 5, 0, 0)
    end_datetime = datetime(2020, 5, 21, 5, 15, 0)
    time_series = pd.to_datetime(time)
    start_index, end_index = np.searchsorted(time_series, [start_datetime, end_datetime])

    rad_refl_t = np.transpose(data[start_index:end_index])
    rad_refl_t[rad_refl_t <= -20] = np.nan
    hgt = np.mean(height[start_index:end_index], axis=0) / 1000
    lat_t = lat[start_index:end_index]
    lon_t = lon[start_index:end_index]
    lat_t[np.isnan(lat_t)] = np.nan
    lon_t[np.isnan(lon_t)] = np.nan

    # filename = "/mnt/datastore/liudddata/result/20200506_droupout/2020051808_predicted_2d_mc.nc"
    # filename = "/mnt/datastore/liudddata/result/20200506_droupout/2020060709_predicted_2d_mc.nc"
    filename = "/mnt/datastore/liudddata/result/20200506_droupout/2020052105_predicted_2d_mc.nc"
    # filename = "/mnt/datastore/liudddata/result/20200506_droupout/2020061306_predicted_2d_mc.nc"
    cth_file = nc.Dataset(filename, 'r')
    cth_data = cth_file.variables['cth'][:]
    cbh_data = cth_file.variables['predicted_mean'][:]
    latitudes, longitudes, cth_values, cbh_values = [], [], [], []
    coord_file_name = 'FY4A_coordinates.nc'
    coord_file_open = nc.Dataset(coord_file_name, 'r')
    lat_fy = coord_file_open.variables['lat'][:, :].T
    lon_fy = coord_file_open.variables['lon'][:, :].T
    lat_fy[np.isnan(lat_fy)] = -90.
    lon_fy[np.isnan(lon_fy)] = 360.
    coord_file_open.close()
    fy_coords = np.deg2rad(np.vstack([lat_fy.ravel(), lon_fy.ravel()]).T)
    fy_coords_cartesian = np.array([np.cos(fy_coords[:, 0]) * np.cos(fy_coords[:, 1]),
                                    np.cos(fy_coords[:, 0]) * np.sin(fy_coords[:, 1]),
                                    np.sin(fy_coords[:, 0])]).T
    is_finite = np.isfinite(fy_coords_cartesian).all(axis=1)
    fy_coords_cartesian_clean = fy_coords_cartesian[is_finite]
    tree = cKDTree(fy_coords_cartesian_clean)
    radius = 2 / 6371
    for lat_val, lon_val in zip(lat_t, lon_t):
        point = np.deg2rad(np.array([lat_val, lon_val]))
        point_cartesian = np.array([np.cos(point[0]) * np.cos(point[1]),
                                    np.cos(point[0]) * np.sin(point[1]),
                                    np.sin(point[0])])
        indices = tree.query_ball_point(point_cartesian, r=radius)
        if indices:
            for idx in indices:
                lat_idx, lon_idx = np.unravel_index(idx, lat_fy.shape)
                latitudes.append(lat_fy[lat_idx, lon_idx])
                longitudes.append(lon_fy[lat_idx, lon_idx])
                cth_values.append(cth_data[lat_idx, lon_idx] * 0.001)
                cbh_values.append(cbh_data[lat_idx, lon_idx] * 0.001)
        else:
            latitudes.append(np.nan);
            longitudes.append(np.nan)
            cth_values.append(np.nan);
            cbh_values.append(np.nan)
    fy_data = {'Latitude': latitudes, 'Longitude': longitudes, 'CTH': cth_values, 'CBH': cbh_values}
    df = pd.DataFrame(fy_data)
    df.dropna(inplace=True)
    df_sorted = df.sort_values(by='Latitude')
    latitude_sorted = df_sorted['Latitude']
    cth_sorted = df_sorted['CTH']
    cbh_sorted = df_sorted['CBH']
    elv /= 1000
    # ...

    # 定义剖面图的纬度范围
    plot_lat_range = (15, 33)
    # plot_lat_range = (-53, -33)
    # plot_lat_range = (-28, -18.5)
    # plot_lat_range = (4, 18)

    fig = plt.figure(dpi=200, figsize=(6, 6))
    ax1 = fig.add_axes([0.25, 0.45, 0.5, 0.5], projection=ccrs.PlateCarree())
    ax2 = fig.add_axes([0.25, 0.2, 0.5, 0.3])

    # 绘制全球地图和轨迹
    set_map(ax1)
    # ------------------ MODIFICATION: Pass latitude range to draw_track function ------------------
    # 在绘制轨迹时传递纬度范围参数，仅绘制该范围内的轨迹
    draw_track(ax1, lon_t, lat_t, lat_range=plot_lat_range)

    # ------------------ MODIFICATION: Call function with custom style ------------------

    start_str = start_datetime.strftime('%Y-%m-%d %H:%M')
    end_str = end_datetime.strftime('%Y-%m-%d %H:%M')
    ax1.set_title(f'From {start_str} to {end_str}', fontsize='small')

    # 绘制剖面图
    draw_cross_section(ax2, lat_t, hgt, rad_refl_t, cth_sorted, cbh_sorted, latitude_sorted, xlims=plot_lat_range)
    draw_elevation(ax2, lat_t, elv[start_index:end_index])

    ax2.set_title('Radar Reflectivity Factor', fontsize='small')

    fig.savefig(
        # '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/2020061306_global_custom_highlight',
        # '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/2020051808_global_custom_highlight',
        # '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/2020060709_global_custom_highlight',
        '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_chart/2020052105_global_custom_highlight',
        bbox_inches='tight')
    plt.show()
    plt.close(fig)